#include <stdio.h>
#include <malloc.h>
#include <math.h>
#include <mpi.h>

void MatrixToCSR(double *pMatrix, double *Value, int* Col, int* RowIndex, int size)
{
	int i, j=0, k=0, l=-1;

	for (i=0; i<size*size; i++)
		if (pMatrix[i]!=0) {

			Value[j]=pMatrix[i];
			Col[j]=(i%size);

			if (l<i/size) RowIndex[k++]=j;

			j++;
			l=i/size;
		}
}

void ProcessInitialization(double** _pMatrix, double** _pVector, double** _pResult, double** _pProcResult, double** _ProcValue, int** _ProcCol, int** _ProcRowIndex, double** _Value, int** _Col, int** _RowIndex, int* _size, int ProcNum, int ProcRank) 
{
	int *SendProcSize, *pSendInd, *SendRI, *SendIndRI;
	int blocksize, intsize, RowNum, RowNum_v, nz=0;

	double *pMatrix, *pVector, *pResult, *pProcResult, *Value, *ProcValue;
	int *Col, *RowIndex, *ProcCol, *ProcRowIndex;
	int size, i, j;

	FILE* fp;

	if (ProcRank==0) {
		/*
		fp = fopen("filename.txt.", "r");
		fscanf(fp, "%d", &size);

		pMatrix=(double *)malloc(size*size*sizeof(double));
		pVector=(double *)malloc(size*sizeof(double));

		for (i=0; i<size*size; ++i) {
			fscanf(fp, "%f", pMatrix[i]);
			if (pMatrix[i]!=0) nz++;
		}

		for (i=size*size; i<size*size+size; i++) 
			fscanf(fp, "%f", pVector[i]);

		fclose(fp);
		*/

	    size=3;
	    MPI_Bcast(&size, 1, MPI_INT, 0, MPI_COMM_WORLD);

	    pMatrix=(double *)malloc(size*size*sizeof(double));
	    pVector=(double *)malloc(size*sizeof(double));

	    for (i=0; i<size; i++) {
		pVector[i]=i+1;

		for (j=0; j<size; j++) {
		    if (abs(i-j)<2) { 
			pMatrix[i*size+j]=i+j+1;
			nz++;
		    }
		    else pMatrix[i*size+j]=0.0;
		}
	    }

	    Value=(double *)malloc(sizeof(double)*nz);
	    Col=(int *)malloc(sizeof(int)*nz);
	    RowIndex=(int*)malloc(sizeof(int)*(size+1));

	    MatrixToCSR(pMatrix, Value, Col, RowIndex, size);

	    ///////////////////////////////////////////////////////
	    RowIndex[size]=nz; //!

	    blocksize=(2*size)/(ProcNum+1); 
	    intsize=size/(ProcNum+1); 

	    MPI_Bcast(pVector, size, MPI_DOUBLE, 0, MPI_COMM_WORLD);

	    SendProcSize=(int *)malloc(sizeof(int)*ProcNum);
	    for (i=0; i<ProcNum-1; i++)
		SendProcSize[i]=RowIndex[i*(blocksize-intsize)+blocksize]-RowIndex[i*(blocksize-intsize)];
	    SendProcSize[ProcNum-1]=RowIndex[size]-RowIndex[(ProcNum-1)*(blocksize-intsize)]; // note
	    MPI_Bcast(SendProcSize, ProcNum, MPI_INT, 0, MPI_COMM_WORLD);

	    pSendInd=(int *)malloc(ProcNum*sizeof(int));
	    pSendInd[0]=0;
	    for (i=0; i<ProcNum-1; i++)
		pSendInd[i+1]=pSendInd[i]+(RowIndex[(i+1)*intsize]-RowIndex[i*intsize]); //?
	    MPI_Bcast(pSendInd, ProcNum, MPI_INT, 0, MPI_COMM_WORLD);
	} 
	else {
	    MPI_Bcast(&size, 1, MPI_INT, 0, MPI_COMM_WORLD);

	    pVector=(double *)malloc(size*sizeof(double));
	    MPI_Bcast(pVector, size, MPI_DOUBLE, 0, MPI_COMM_WORLD);

	    SendProcSize=(int *)malloc(sizeof(int)*ProcNum);
	    MPI_Bcast(SendProcSize, ProcNum, MPI_INT, 0, MPI_COMM_WORLD);

	    pSendInd=(int *)malloc(ProcNum*sizeof(int));
	    MPI_Bcast(pSendInd, ProcNum, MPI_INT, 0, MPI_COMM_WORLD);

	    blocksize=(2*size)/(ProcNum+1); 
	    intsize=size/(ProcNum+1);
	}

	pResult=(double *)malloc(size*sizeof(double));


	if (ProcRank!=ProcNum-1) {
	    RowNum=blocksize;
	    RowNum_v=intsize;
	}
	else {
	    RowNum=size-intsize*(ProcNum-1)+1; //!
	    RowNum_v=RowNum;
	}

	ProcValue=(double *)malloc(sizeof(double)*SendProcSize[ProcRank]);
	ProcRowIndex=(int *)malloc(sizeof(int)*RowNum);
	ProcCol=(int *)malloc(sizeof(int)*SendProcSize[ProcRank]);
	pProcResult=(double *)malloc(sizeof(double)*RowNum_v);

	SendRI=(int *)malloc(sizeof(int)*ProcNum);
	SendIndRI=(int *)malloc(sizeof(int)*ProcNum);
	
	for (i=0; i<ProcNum-1; i++)
	    SendRI[i]=blocksize;
	SendRI[ProcNum-1]=size-intsize*(ProcNum-1)+1;

	SendIndRI[0]=0;
	for (i=0; i<ProcNum-1; i++)
	    SendIndRI[i+1]=SendIndRI[i]+intsize;

	MPI_Scatterv(Value, SendProcSize, pSendInd, MPI_DOUBLE, ProcValue, SendProcSize[ProcRank], MPI_DOUBLE, 0, MPI_COMM_WORLD);
	MPI_Scatterv(Col, SendProcSize, pSendInd, MPI_INT, ProcCol, SendProcSize[ProcRank], MPI_INT, 0, MPI_COMM_WORLD);
	MPI_Scatterv(RowIndex, SendRI, SendIndRI, MPI_INT, ProcRowIndex, SendRI[ProcRank], MPI_INT, 0, MPI_COMM_WORLD);

	*_size=size;
	*_pMatrix=pMatrix;
	*_pVector=pVector;
	*_pResult=pResult;
	*_pProcResult=pProcResult;
	*_Value=Value;
	*_ProcValue=ProcValue;
	*_Col=Col;
	*_RowIndex=RowIndex;
	*_ProcCol=ProcCol;
	*_ProcRowIndex=ProcRowIndex;

	free(SendProcSize);
	free(SendRI);
	free(pSendInd);
	free(SendIndRI);
}

void ParallelResultCalculation(double *ProcValue, int* ProcCol, int* ProcRowIndex, double *pVector, double *pProcResult, int size, int ProcNum, int ProcRank)
{
	int i, j, i1, j1, j2;
	double sum;

	int blocksize=2*size/(ProcNum+1);
	int intsize=size/(ProcNum+1);

	if (ProcRank != ProcNum-1) i1=intsize;
	else i1=size-intsize*(ProcNum-1);

//	See U

	for (i=0; i<i1; i++) {
		sum=0.0;

		for (j=ProcRowIndex[i]; j<ProcRowIndex[i+1]; j++)
			sum+=ProcValue[j-ProcRowIndex[0]]*pVector[ProcCol[j-ProcRowIndex[0]]];
		pProcResult[i]=sum;
	}
}

void ResultReplication(double *pProcResult, double *pResult, int size, int ProcNum, int ProcRank)
{
	int *pReceiveNum; 
	int *pReceiveInd; 
	int i;
	int blocksize=2*size/(ProcNum+1);
	int intsize=size/(ProcNum+1);

	pReceiveNum=(int *)malloc(ProcNum*sizeof(int));
	pReceiveInd=(int *)malloc(ProcNum*sizeof(int));

	pReceiveInd[0]=0;
	for (i=0; i<ProcNum-1; i++) {
	    pReceiveNum[i]=intsize;
	    pReceiveInd[i+1]=pReceiveInd[i]+pReceiveNum[i];
	}
	pReceiveNum[ProcNum-1]=size-(ProcNum-1)*intsize;

	MPI_Allgatherv(pProcResult, pReceiveNum[ProcRank], MPI_DOUBLE, pResult, pReceiveNum, pReceiveInd, MPI_DOUBLE, MPI_COMM_WORLD);

	free(pReceiveNum);
	free(pReceiveInd);
}

int main(int argc, char *argv[])
{
	int size;

	double *pMatrix, *pVector, *pResult, *Value;
	int *RowIndex, *Col;

	int ProcNum, ProcRank;

	double *pProcRows, *pProcResult, *ProcValue;
	int *ProcRowIndex, *ProcCol;

	double start, finish;
	int i,j;

	MPI_Init(&argc, &argv);

	MPI_Comm_size(MPI_COMM_WORLD, &ProcNum);
	MPI_Comm_rank(MPI_COMM_WORLD, &ProcRank);

	start=MPI_Wtime();

	ProcessInitialization(&pMatrix, &pVector, &pResult, &pProcResult, &ProcValue, &ProcCol, &ProcRowIndex, &Value, &Col, &RowIndex, &size, ProcNum, ProcRank);

	ParallelResultCalculation(ProcValue, ProcCol, ProcRowIndex, pVector, pProcResult, size, ProcNum, ProcRank);

	ResultReplication(pProcResult, pResult, size, ProcNum, ProcRank);

	finish=MPI_Wtime();

	MPI_Finalize();

	/*
	for (i=0; i<size; i++)
		printf("%f\n", pResult[i]);
	*/

	//printf("\ntime: %f\n", finish-start);

	if (ProcRank==0) { free(pMatrix); free(Value); free(RowIndex); free(Col); }
	free(pVector);
	free(pResult);
	free(ProcValue);
	free(ProcRowIndex);
	free(ProcCol);
	free(pProcResult);

	return 0;
}
